import arviz as az
import numpy as np
from arviz.labels import MapLabeller

print(az.__version__)

datadict = {
    "a": np.random.randn(1, 100),
    "b": np.random.randn(1, 100),
    "c": np.random.randn(1, 100, 2),
    "d": np.random.randn(1, 100, 2),
    "x": np.random.randn(1, 100, 2, 1),
    "y": np.random.randn(1, 100, 2, 1),
}
trace = az.convert_to_inference_data(datadict)

reference_values = {
    "a": 0.0,
    "b": 0.0,
    "c": [-0.5, 0.5],
    "d": [-0.5, 0.5],
    "x": np.array([[-0.5], [0.5]]),
    "y": np.array([[-0.5], [0.5]]),
}

var_name_map = {
    "a": r"$\alpha$ ($\mu$m)",
    "b": r"$\beta$ ($\mu$m)",
    "x": "x (cm)",
    "y": "y (cm)",
}
labeller = MapLabeller(var_name_map)

az.plot_pair(
    trace,
    var_names=["a", "b"],
    reference_values=reference_values,
    labeller=labeller,
    show=True,
)
az.plot_pair(
    trace,
    var_names=["c", "d"],
    reference_values=reference_values,
    labeller=labeller,
    show=True,
)
az.plot_pair(
    trace,
    var_names=["x", "y"],
    reference_values=reference_values,
    labeller=labeller,
    combine_dims={"x_dim_0", "y_dim_0"},
    show=True,
)